% 速度Verlet法计算多体系统运动！
% 使用AI辅助
% Gitee Repo

clc
clear

global N

L = 5;
N = 10; % 粒子个数

% 添加摩擦力
allowFriction = 0;
% 添加刚性边界，粒子撞击边界后将反弹
allowBoxed = 1;
% 启用Nose-Hoover型温度（总动能）控制算法
allowKEControl = 1;
zeta = 0; % 热浴变量
Q = 50; % 热浴系数
KE_target = 5; % 目标总动能

% 粒子位置 N*3
% 行：粒子索引，列：x,y,z
r0 = 4*rand(N,3)-2;
% 粒子速度 N*3
v0 = 2*rand(N,3)-1;

% 粒子位置、速度居中
% 根据动量守恒，如果所有粒子速度和不为0，那么系统会整体运动离开中心
r0 -= sum(r0,1)/N;
v0 -= sum(v0,1)/N;

r1 = zeros(N,3);
v1 = zeros(N,3);

dt = 0.01;

% 计算两个粒子之间的相互作用力
function F = compute_F(r1,r2)
    % r1, r2: position of particles.
    Rvec = r2-r1;

    R2 = sum(Rvec.^2);
    R = sqrt(R2);

    % LJ 型力，在两个粒子距离较远时吸引，在较近时排斥
    F = (+0.2/R.^2-1/R.^6)*Rvec;
end

%  计算两两粒子相互作用力
function Fmat = compute_interaction(r)
    global N
    Fmat = zeros(N,N,3);
    for i = 1:N
        for j = i+1:N
            F = compute_F(r(i,:),r(j,:));
            Fmat(i,j,:) = F(:);
        end
    end
    % 牛三
    Fmat += -permute(Fmat,[2,1,3]);
end

figure()
Fmat = [];
KE = 0;
for tick = 1:10000

    % 计算半步长速度
    Fmat = compute_interaction(r0);
    vhalf = v0 + 1/2*squeeze(sum(Fmat,2))*dt - 1/2*zeta*v0*dt;

    % 计算下一时刻位置
    r1 = r0 + vhalf*dt;

    % 计算下一时刻速度
    Fmat = compute_interaction(r1);
    v1 = vhalf + 1/2*squeeze(sum(Fmat,2))*dt - 1/2*zeta*vhalf*dt;

    % 更新位置与速度
    r0 = r1;
    v0 = v1;

    % “摩擦”
    % 在有摩擦的情况下，应该能够观察到粒子最终形成比较稳定的结构
    if allowFriction
        v0*=0.998;
    end

    % 刚性边界，粒子撞击边界后将反弹
    if allowBoxed
        %定义辅助逻辑索引
        belowLowerBound = r0 < -L;
        aboveUpperBound = r0 > L;

        %应用边界条件到位置r0
        r0(belowLowerBound) = -L;
        r0(aboveUpperBound) = L;

        %反转越界的速度分量
        v0(belowLowerBound) = -v0(belowLowerBound);
        v0(aboveUpperBound) = -v0(aboveUpperBound);
    end

    % 总动能控制
    KE = 1/2*sum(v1(:).^2);
    if allowKEControl
        zeta += (KE-KE_target)/Q*dt;
    else
        zeta = 0;
    end

    % 判断迭代是否发散
    if any(isnan(r0)) || any(isinf(r0))
        error('diverge');
    end

    if all(abs(v0)<0.01)
        printf('All particles have stopped \n')
        break
    end

    % 画图
    if mod(tick,10)==0
        clf
        hold on
        axis equal
        axis([-L L -L L -L L])

        for i = 1:N
            scatter3(r1(i,1),r1(i,2),r1(i,3));
        end
        view([30,45])
        drawnow
        pause(0.02)

        disp(tick)
        disp(KE)
    end
end
